import torch
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from .models import (
    ModelBiLSTM, 
)
from .dataloader import (
    SignalFeaData1s,
    generate_offsets,
)
from .utils.process_utils import count_line_num
import argparse

def main():
    parser = argparse.ArgumentParser(description="Training script for ModelBiLSTM using PyTorch Lightning")
    parser.add_argument('--train_file', type=str, required=True)
    parser.add_argument('--valid_file', type=str, required=True)
    parser.add_argument('--model_dir', type=str, required=True)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--max_epochs', type=int, default=20)
    parser.add_argument('--gpus', type=int, default=1, help="Number of GPUs to use.")
    parser.add_argument('--init_model', type=str, default=None, help="Path to load the pre-trained model")
    args = parser.parse_args()

    # Dataset setup
    train_dataset = SignalFeaData1s(args.train_file, generate_offsets(args.train_file), count_line_num(args.train_file))
    val_dataset = SignalFeaData1s(args.valid_file, generate_offsets(args.valid_file), count_line_num(args.valid_file))

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)

    # Model setup
    model = ModelBiLSTM()
    if args.init_model:
        model = ModelBiLSTM.load_from_checkpoint(args.init_model)

    # Callbacks for saving checkpoints and early stopping
    checkpoint_callback = ModelCheckpoint(
        dirpath=args.model_dir,
        filename='model-{epoch:02d}-{val_loss:.2f}',
        save_top_k=3,
        monitor='val_loss',
        mode='min'
    )
    early_stopping = EarlyStopping(
        monitor='val_loss',
        patience=5,
        mode='min'
    )

    # Training setup with Lightning
    trainer = pl.Trainer(
        max_epochs=args.max_epochs,
        gpus=args.gpus,
        accelerator='gpu',
        callbacks=[checkpoint_callback, early_stopping]
    )
    trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)

if __name__ == "__main__":
    main()
